from __future__ import print_function, division
import time
#Importing relevant modules and packages
import numpy as np
import matplotlib.pyplot as plt
import random
from RichardsModel_1D import *
from Parameters import *
from scipy import integrate
from tqdm import tqdm



p=Loam() #Soil parameters

#Space and time parameters
Nr,Nt,Nz,dr,dt,dz,Np,NN,Nx,Nw,Ny,Nv,Hz=circular_parameters()
DeltaT,Nsim,Tsim=time_parameters()

N = 100 # Number of particles

#Irrigation amount
uu=np.zeros(DeltaT)
# for i in range(30):
#     uu[40*i]=-3e-06 #Moisture is supplied only at the first time instant
    # uu[0*i]=-1e-05



Kc=0.88*np.ones(Nsim)      #Crop coefficient
Kc1=0.88*np.ones(Tsim)      #Crop coefficient
Et=1.389e-08*np.ones(Nsim) #Reference evapotranspiration
Et1=1.389e-08*np.ones(DeltaT) #Reference evapotranspiration



#Initial condition [pressure head]
x0=-0.5*np.ones(Nz)

#Initial particle filter state
p0=1.02*x0

#Numpy array to store the simulation results [pressure head]
xx=np.zeros((Tsim+1,Nz))
xx1=np.zeros((Nsim+1,Nz))
ode1=np.zeros((Tsim+1,Nz))
xx[0,:]=x0
xx1[0,:]=x0

# array for filtering
x_out = np.zeros((Nsim+1,Nz))
x_est = np.zeros((Nsim+1,Nz))
x_est_out = np.zeros((Nsim+1,Nz))
x_out[0,:] = x0  #the actual output vector for measurement values.
x_est[0,:] = p0 # time by time output of the particle filters estimate
x_est_out[0,:] = p0 # the vector of particle filter estimates.

# array for particles
P=np.zeros((N,Nz)) # Dimension of particles
x_P_update=np.zeros((N,Nz))
Z_update = np.zeros((N,Nz))
w = np.zeros(N)
xnew = np.zeros((N,Nz))

#Numpy array to store the simulation results [volumetric moisture content]
yy=np.zeros((Nsim+1,Nz))
yy1=np.zeros((Nsim+1,Nz))
yy[0,:]=getOutputs1D_allnodes(x0,p).ravel()
yy1[0,:]=getOutputs1D_allnodes(x0,p).ravel()


#Model uncertainty

np.random.seed(1)
ww=0.0000*np.random.randn(DeltaT,Nz) #If no model uncertainty is present, ww is multiplied by 0
np.random.seed(2)
www=0.00000000*np.random.randn(Nsim,Nz)
np.random.seed(3)
vv=0.1*np.random.randn(Nsim,Nz) #measurement noise
np.random.seed(4)
vvv=0.0*np.random.randn(N,Nz) # initial particle covariance
# ww1=0.000001
# ww=0*np.random.normal(0,np.sqrt(ww1),(Nsim,Nz))
# www=np.random.normal(0,np.sqrt(ww1),(Nsim,Nz))
# vv1=0.01
# vv=np.random.normal(0,np.sqrt(vv1),(Nsim,Nz))
# vvv1=0.000001
# vvv=np.random.normal(0,np.sqrt(vvv1),(Nsim,Nz))

#unknown inputs
a1=0*np.ones((DeltaT+1,Nz))
a=0.00000005*np.ones((Nsim+1,Nz))
a_est=np.zeros((Nsim+1,Nz))
a_est[0,:]=0.0000000

gamma=[0.3,0.3,0.3,0.3,0.3,0.3,0.3,0.3,0.3,0.3,0.3,0.1,0.1,0.1,0.02,0.005]


def ode(x,t,u,w,ui,Kc,ET0,p):  
    return RichardsModel_1D(x,u,w,ui,Kc,ET0,p)

 

e0 = time.time()
# NOW THE FUNCTION IS COMPILED, RE-TIME IT EXECUTING FROM CACHE

# initialize particle filter
for j in range (1, N+1):
    P[j-1,:]=p0+vvv[j-1,:]
    
for i in range(1, Nsim+1):
    # Generate observed data (x bias)
    1


for i in tqdm(range(1, Nsim+1)):
    ## Use Scipy.integrate.odeint method
    # sol = integrate.odeint(ode, xx[i-1], [0,DeltaT], args=(uu[i-1],ww[i-1],Kc[i-1],Et[i-1],p))
    # xk1 = sol[1,:]   
    
    ## Use Scipy.integrate.solve_ivp method
    # Method options: 
    # 'RK45': Explicit Runge-Kutta method of order 5(4)		4.804461243
    # 'RK23': Explicit Runge-Kutta method of order 3(2) 	2.625176191		
    # 'DOP853': Explicit Runge-Kutta method of order 8	    8.995131016
    # 'Radau': Implicit Runge-Kutta method of the Radau IIA family of order 5  8.726206541		
    # 'BDF': Implicit multi-step variable-order (1 to 5) method based on a backward differentiation formula for the derivative approximation   6.288985729		
    # 'LSODA': Adams/BDF method with automatic stiffness detection and switching. This is a wrapper of the Fortran solver from ODEPACK	   1.380549431
    
    # Generate observed data (y bias)
    # sol_1 = integrate.solve_ivp(fun=lambda t, y: ode(y, t, uu[i-1],www[i-1],a1[i-1],Kc[i-1],Et[i-1],p),t_span=[0,DeltaT],y0=tuple(xx1[i-1]),method='LSODA')
    # xk2 = sol_1.y[:,-1]   
    # xx1[i,:]=xk2
    # yk2=getOutputs1D_allnodes(xk2,p).ravel()
    # yy1[i,:]=yk2+vv[i-1,:]+a[i-1,:]
    
     # REM-PF
    for j in range (1,N+1):
        sol_2 = integrate.solve_ivp(fun=lambda t, y: ode(y, t, uu[i-1],www[i-1],a_est[i-1],Kc[i-1],Et[i-1],p),t_span=[0,DeltaT],y0=tuple(P[j-1]),method='LSODA')
        pp=sol_2.y[:,-1]
        x_P_update[j-1,:]=pp
        zz=getOutputs1D_allnodes(pp,p).ravel()
        Z_update[j-1,:]=zz
        dist=np.linalg.norm(Z_update[j-1,:]-yy1[i,:],ord=2)
        w[j-1]=np.exp(-((dist)**2/(2*0.1)))
    w=w/np.sum(w)
    
    c=[0]*N
    c[0]=w[0]
    for j in range(N):
        c[j]=c[j-1]+w[j]
        
    for j in range(N):  
        d=random.uniform(0,1)
        for k in range(N):
            if d <= c[k]:
                xnew[j,:]=x_P_update[k,:]
                break
    x_est[i,:]=np.sum(xnew,axis=0)/N     
    
    # Save data in arrays for later plotting
    x_out[i,:] = xx1[i,:]
    x_est_out[i,:] = x_est[i,:]
    ode1[i-1,:]= RichardsModel_1D(tuple(x_est_out[i-1]), uu[i-1],ww[i-1],a1[i-1],Kc[i-1],Et[i-1],p)
    # a_est[i,:]=(np.ones(16)-gamma)*a_est[i-1,:]+gamma*((x_est_out[i,:]-x_est_out[i-1,:])/DeltaT-ode1[i-1,:])

    P=xnew
 

elapsed_time=time.time()-e0

print("Model (1D) runs successfully after",elapsed_time,"seconds")


t= list(range(Nsim+1))
# t= list(range(N))
# plt.subplot(3, 1, 1)    
# # plt.plot(t,P[:,1],'r')
# plt.plot(t,xx1[:,0],'r',t,xx[:,0],'b')
plt.figure(1)
for i in range (Nz):
    plt.subplot(4, 4, i+1)
# plt.plot(t,x_out[:,1],'r')
    plt.plot(t,x_out[:,i],'r',t,x_est_out[:,i],'b')

plt.figure(2)
for i in range (Nz):
    plt.subplot(4, 4, i+1)
    plt.plot(t,a_est[:,i],'b',t,a[:,i],'r')
plt.show()